#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""代码保护模块
负责对关键代码进行保护和混淆"""

import os
import sys
import types
import marshal
import base64
import zlib
import importlib
import importlib.abc
import importlib.util
import inspect
import threading
from utils import Logger
from utils.security_utils import SecurityUtils

class CodeProtector:
    """代码保护者
    负责对关键代码进行保护和混淆"""
    
    def __init__(self):
        """初始化代码保护者"""
        self.security_utils = SecurityUtils()
        self._protected_modules = {}
        self._protection_lock = threading.Lock()
        Logger.log_info("代码保护者初始化完成")
    
    def protect_module(self, module_name, module_path=None):
        """保护指定的模块
        
        Args:
            module_name: 模块名称
            module_path: 模块路径，如果为None则尝试从已导入的模块中查找
            
        Returns:
            bool: 是否保护成功
        """
        with self._protection_lock:
            Logger.log_info(f"开始保护模块: {module_name}")
            
            try:
                # 获取模块对象
                module = None
                
                if module_path:
                    # 从文件加载模块
                    module = self._load_module_from_file(module_name, module_path)
                else:
                    # 尝试从已导入的模块中获取
                    if module_name in sys.modules:
                        module = sys.modules[module_name]
                    else:
                        Logger.log_error(f"模块 {module_name} 未找到")
                        return False
                
                # 检查模块是否有效
                if not module:
                    Logger.log_error(f"无法加载模块: {module_name}")
                    return False
                
                # 保护模块内容
                protected_code = self._protect_module_code(module)
                
                if protected_code:
                    # 存储保护后的模块
                    self._protected_modules[module_name] = protected_code
                    Logger.log_info(f"模块保护成功: {module_name}")
                    return True
                else:
                    Logger.log_error(f"模块保护失败: {module_name}")
                    return False
            except Exception as e:
                Logger.log_error(f"保护模块 {module_name} 时发生错误: {str(e)}")
                return False
    
    def _load_module_from_file(self, module_name, module_path):
        """从文件加载模块
        
        Args:
            module_name: 模块名称
            module_path: 模块文件路径
            
        Returns:
            module: 加载的模块对象
        """
        try:
            # 检查文件是否存在
            if not os.path.exists(module_path):
                Logger.log_error(f"模块文件不存在: {module_path}")
                return None
            
            # 使用importlib加载模块
            spec = importlib.util.spec_from_file_location(module_name, module_path)
            if not spec:
                Logger.log_error(f"无法为模块 {module_name} 创建spec")
                return None
            
            module = importlib.util.module_from_spec(spec)
            sys.modules[module_name] = module
            
            try:
                spec.loader.exec_module(module)
                return module
            except Exception as e:
                Logger.log_error(f"执行模块 {module_name} 时发生错误: {str(e)}")
                # 从sys.modules中移除
                if module_name in sys.modules:
                    del sys.modules[module_name]
                return None
        except Exception as e:
            Logger.log_error(f"从文件加载模块 {module_name} 时发生错误: {str(e)}")
            return None
    
    def _protect_module_code(self, module):
        """保护模块代码
        
        Args:
            module: 模块对象
            
        Returns:
            str: 保护后的代码
        """
        try:
            # 获取模块中的所有函数
            functions = []
            
            for name, obj in inspect.getmembers(module):
                if inspect.isfunction(obj) or inspect.ismethod(obj):
                    functions.append((name, obj))
            
            # 保护关键函数
            for name, func in functions:
                # 检查是否是我们需要保护的函数
                if self._should_protect_function(name, func):
                    protected_func = self._protect_function(func)
                    # 替换原函数
                    setattr(module, name, protected_func)
            
            # 对模块进行整体保护
            # 在这个示例中，我们只返回模块名称作为标记
            # 实际应用中，这里应该有更复杂的保护逻辑
            return f"PROTECTED_MODULE:{module.__name__}"
        except Exception as e:
            Logger.log_error(f"保护模块代码时发生错误: {str(e)}")
            return None
    
    def _should_protect_function(self, function_name, function_obj):
        """判断函数是否需要被保护
        
        Args:
            function_name: 函数名称
            function_obj: 函数对象
            
        Returns:
            bool: 是否需要保护
        """
        # 这里可以定义需要保护的函数的判断逻辑
        # 在这个示例中，我们保护所有非私有函数
        if not function_name.startswith('_'):
            # 检查函数是否有docstring，并且包含敏感标记
            if function_obj.__doc__ and ("private" in function_obj.__doc__.lower() or "sensitive" in function_obj.__doc__.lower()):
                return True
            
            # 检查函数是否在特定模块中
            if function_obj.__module__ and "security" in function_obj.__module__.lower():
                return True
        
        return False
    
    def _protect_function(self, func):
        """保护函数
        
        Args:
            func: 要保护的函数
            
        Returns:
            function: 保护后的函数
        """
        try:
            # 获取函数的源代码
            source = inspect.getsource(func)
            
            # 在实际应用中，这里应该对源代码进行混淆和保护
            # 在这个示例中，我们只是简单地记录函数被保护了
            Logger.log_info(f"保护函数: {func.__name__} 在 {func.__module__}")
            
            # 返回原始函数
            # 在实际应用中，这里应该返回保护后的函数
            return func
        except Exception as e:
            Logger.log_error(f"保护函数 {func.__name__} 时发生错误: {str(e)}")
            # 如果发生错误，返回原始函数
            return func
    
    def obfuscate_code(self, code_string):
        """混淆代码字符串
        
        Args:
            code_string: 原始代码字符串
            
        Returns:
            str: 混淆后的代码字符串
        """
        try:
            Logger.log_info("执行代码混淆...")
            
            # 在这个示例中，我们简单地使用base64和zlib进行混淆
            # 在实际应用中，这里应该使用更复杂的混淆技术
            
            # 压缩代码
            compressed = zlib.compress(code_string.encode('utf-8'))
            
            # 进行base64编码
            encoded = base64.b64encode(compressed)
            
            # 构建执行混淆代码的代码
            obfuscated_code = f"""
import zlib
import base64
exec(zlib.decompress(base64.b64decode({encoded})))
"""
            
            return obfuscated_code
        except Exception as e:
            Logger.log_error(f"代码混淆时发生错误: {str(e)}")
            # 如果混淆失败，返回原始代码
            return code_string
    
    def create_protected_executor(self, code_string):
        """创建一个受保护的执行器
        
        Args:
            code_string: 要执行的代码字符串
            
        Returns:
            function: 可以执行受保护代码的函数
        """
        try:
            Logger.log_info("创建受保护的执行器...")
            
            # 混淆代码
            obfuscated_code = self.obfuscate_code(code_string)
            
            # 创建一个函数来执行混淆后的代码
            def protected_executor(*args, **kwargs):
                """受保护的代码执行器"""
                # 在实际应用中，这里应该有更复杂的安全检查和执行逻辑
                try:
                    # 定义一个局部环境来执行代码
                    local_env = {
                        'args': args,
                        'kwargs': kwargs
                    }
                    
                    # 执行混淆后的代码
                    exec(obfuscated_code, local_env)
                    
                    # 返回执行结果（如果有）
                    return local_env.get('result')
                except Exception as e:
                    Logger.log_error(f"执行受保护代码时发生错误: {str(e)}")
                    return None
            
            return protected_executor
        except Exception as e:
            Logger.log_error(f"创建受保护执行器时发生错误: {str(e)}")
            # 如果创建失败，返回一个简单的执行器
            def simple_executor(*args, **kwargs):
                try:
                    exec(code_string)
                except Exception:
                    pass
                return None
            
            return simple_executor
    
    def get_protected_modules(self):
        """获取所有受保护的模块
        
        Returns:
            dict: 受保护模块的字典
        """
        return self._protected_modules.copy()

# 示例用法
if __name__ == "__main__":
    protector = CodeProtector()
    
    # 示例函数
    def sensitive_function():
        """这是一个敏感函数，包含private和sensitive信息"""
        return "敏感数据"
    
    # 混淆一些代码
    sample_code = """
# 这是一个简单的示例代码
result = 'Hello, World!'
print(result)
"""
    
    obfuscated_code = protector.obfuscate_code(sample_code)
    print("混淆后的代码:")
    print(obfuscated_code)
    
    # 创建一个受保护的执行器
    executor = protector.create_protected_executor(sample_code)
    print("\n执行受保护的代码:")
    result = executor()
    print(f"执行结果: {result}")